package ssh

import (
	"errors"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"time"
)

type Option struct {
	Addr        string
	User        string
	Password    string //类型为password
	Type        string //password 或者 key
	Timeout     time.Duration
	key         string //类型为key
	keyPassword string //类型为key
	High        int    //高
	Wide        int    //宽
}

func (this *Option) new() *Option {
	if len(this.Addr) == 0 {
		this.Addr = "192.168.10.44"
	}
	if len(this.User) == 0 {
		this.User = "root"
	}
	if this.Timeout == 0 {
		this.Timeout = time.Second
	}
	if this.High == 0 {
		this.High = 32
	}
	if this.Wide == 0 {
		this.Wide = 300
	}
	return this
}

func NewClient(op *Option) (*Client, error) {
	op.new()
	config := &ssh.ClientConfig{
		Timeout:         op.Timeout,
		User:            op.User,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Auth:            []ssh.AuthMethod{ssh.Password(op.Password)},
	}
	switch op.Type {
	case "key":
		signer, err := ssh.ParsePrivateKeyWithPassphrase([]byte(op.key), []byte(op.keyPassword))
		if err != nil {
			return nil, err
		}
		config.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
	}
	sshClient, err := ssh.Dial("tcp", op.Addr, config)
	if err != nil {
		return nil, err
	}
	session, err := sshClient.NewSession()
	if err != nil {
		return nil, err
	}
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,     // 禁用回显（0禁用，1启动）
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, //output speed = 14.4kbaud
	}
	output, err := session.StdoutPipe()
	if err != nil {
		return nil, err
	}
	outputErr, err := session.StderrPipe()
	if err != nil {
		return nil, err
	}
	writer, err := session.StdinPipe()
	if err != nil {
		return nil, err
	}
	if err := session.RequestPty("xterm", op.High, op.Wide, modes); err != nil {
		return nil, err
	}
	if err := session.Shell(); err != nil {
		return nil, err
	}
	return &Client{
		Session:   session,
		op:        op,
		writer:    writer,
		output:    output,
		err:       outputErr,
		closeSign: make(chan byte, 1),
		running:   false,
	}, nil
}

func MustClient(op *Option) *Client {
	t := time.Second
again:
	s, err := NewClient(op)
	if err != nil {
		if t < time.Second*30 {
			t = 2 * t
		}
		log.Printf("[错误][%s] %s,等待%s秒重试", op.Addr, err.Error(), t)
		time.Sleep(t)
		goto again
	}
	return s
}

func RedialClient(op *Option, fn ...func(*Client)) *Client {
	return MustClient(op).Redial(fn...)
}

type Client struct {
	*ssh.Session
	op        *Option
	writer    io.Writer
	output    io.Reader
	err       io.Reader
	closeSign chan byte
	closeErr  error
	running   bool
	dealFunc  func(*Client, []byte)
	closeFunc func(*Client)
	closeAll  bool
}

func (this *Client) CloseTest() {
	this.close(errors.New("主动关闭"))
}

func (this *Client) Close() {
	this.closeAll = true
	this.close(errors.New("主动关闭"))
}

func (this *Client) SetDealFunc(fn func(*Client, []byte)) {
	this.dealFunc = fn
}

func (this *Client) SetCloseFunc(fn func(*Client)) {
	this.closeFunc = fn
}

func (this *Client) close(err error) {
	if err != nil {
		if this.closeErr == nil {
			this.closeErr = err
			close(this.closeSign)
			this.Session.Close()
			if this.closeFunc != nil {
				this.closeFunc(this)
			}
		}
	}
}

// Redial 重连
func (this *Client) Redial(fn ...func(*Client)) *Client {
	this.SetCloseFunc(func(*Client) {
		if this.closeAll {
			log.Printf("[信息][%s] 连接断开", this.op.Addr)
			return
		}
		log.Printf("[信息][%s] 连接断开,尝试重新连接\n", this.op.Addr)
		time.Sleep(time.Second)
		*this = *MustClient(this.op)
		this.Redial(fn...)
	})
	for _, v := range fn {
		v(this)
	}
	go this.Run()
	return this
}

// SetWindow 修改窗口大小
func (this *Client) SetWindow(h, w int) error {
	return this.Session.WindowChange(h, w)
}

// Write 写入数据
func (this *Client) Write(cmd string) error {
	if !this.running {
		go this.Run()
	}
	bs := []byte(cmd + "\n")
	if _, err := this.writer.Write(bs); err != nil {
		this.close(err)
		return err
	}
	return nil
}

func (this *Client) Run() error {
	return this.Wait()
}

func (this *Client) Wait() error {
	if this.running {
		select {
		case <-this.closeSign:
			return this.closeErr
		}
	}
	this.running = true
	go func() {
		for {
			time.Sleep(time.Millisecond * 200)
			buf := make([]byte, 512)
			n, err := this.output.Read(buf)
			if err != nil {
				this.close(err)
				return
			}
			if this.dealFunc != nil {
				this.dealFunc(this, buf[:n])
			} else {
				fmt.Println(string(buf[:n]))
			}
		}
	}()
	return this.Session.Wait()
}
